// Usage sudo ./rx local_ip <device> (device is optional)
// e.g. sudo ./rx 192.168.1.2

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <ctype.h>
#include <time.h>
#include <pcap.h>

#define __FAVOR_BSD

#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <arpa/inet.h>
#include <signal.h>

/* default snap length (maximum bytes per packet to capture) */
#define SNAP_LEN 1518

/* ethernet headers are always exactly 14 bytes [1] */
#define SIZE_ETHERNET 14

/* Ethernet addresses are 6 bytes */
#define ETHER_ADDR_LEN  6

/* Ethernet header */
struct sniff_ethernet {
  u_char  ether_dhost[ETHER_ADDR_LEN];    /* destination host address */
  u_char  ether_shost[ETHER_ADDR_LEN];    /* source host address */
  u_short ether_type;                     /* IP? ARP? RARP? etc */
};

/* IP header */
struct sniff_ip {
  u_char  ip_vhl;                 /* version << 4 | header length >> 2 */
  u_char  ip_tos;                 /* type of service */
  u_short ip_len;                 /* total length */
  u_short ip_id;                  /* identification */
  u_short ip_off;                 /* fragment offset field */
  #define IP_RF 0x8000            /* reserved fragment flag */
  #define IP_DF 0x4000            /* dont fragment flag */
  #define IP_MF 0x2000            /* more fragments flag */
  #define IP_OFFMASK 0x1fff       /* mask for fragmenting bits */
  u_char  ip_ttl;                 /* time to live */
  u_char  ip_p;                   /* protocol */
  u_short ip_sum;                 /* checksum */
  struct  in_addr ip_src,ip_dst;  /* source and dest address */
};
#define IP_HL(ip)               (((ip)->ip_vhl) & 0x0f)
#define IP_V(ip)                (((ip)->ip_vhl) >> 4)

/* TCP header */
typedef u_int tcp_seq;

struct sniff_tcp {
  u_short th_sport;               /* source port */
  u_short th_dport;               /* destination port */
  tcp_seq th_seq;                 /* sequence number */
  tcp_seq th_ack;                 /* acknowledgement number */
  u_char  th_offx2;               /* data offset, rsvd */
#define TH_OFF(th)      (((th)->th_offx2 & 0xf0) >> 4)
  u_char  th_flags;
  #define TH_FIN  0x01
  #define TH_SYN  0x02
  #define TH_RST  0x04
  #define TH_PUSH 0x08
  #define TH_ACK  0x10
  #define TH_URG  0x20
  #define TH_ECE  0x40
  #define TH_CWR  0x80
  #define TH_FLAGS        (TH_FIN|TH_SYN|TH_RST|TH_ACK|TH_URG|TH_ECE|TH_CWR)
  u_short th_win;                 /* window */
  u_short th_sum;                 /* checksum */
  u_short th_urp;                 /* urgent pointer */
};

long long ns_offset;
long long get_time_ns() {
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &ts);
  long long time = ((long long)ts.tv_sec) * 1000000000 + ts.tv_nsec;
  return time;
}

#define READ_IP 0
#define READ_OFFSET 1
#define READ_TX_TIME 2
#define READ_ID 3
#define DONE 4

// Packet format: planetlab_node_ip,(float)offset_sec,(long long)time_ns
void print_payload(const u_char *payload, int len, long long time_received_ns)
{
	char buffer[24];
  char node_ip[16];
	long long time_transmitted_ns;
	float offset_sec;
	long long id;	
  int i;
	int error = 0;
	int start = 0;
	int state = READ_IP;
	memset(node_ip, 0, sizeof(node_ip));
	memset(buffer, 0, sizeof(buffer));
  
	// Copy comma-delimited portions of payload to a buffer
	// ending with a null byte so we can safely use scanf
  for (i = 0; ; i++) {
		if (payload[i] == 0) {
			//fprintf(stderr, "  no more input (null byte)\n");
			break;
		}
		if (i >= len) {
			//fprintf(stderr, "  no more input (len)\n");
			break;
		}
		if (i - start >= sizeof(buffer)) {
			//fprintf(stderr, "  buffer length exceeded\n");
			break;
		}

    if (payload[i] != ',') {
			buffer[i - start] = payload[i];
			continue;
		}

		buffer[i] = '\0';
		//fprintf(stdout, "  comma, i=%d, start=%d, buffer=%s\n", i, start, buffer);
		switch (state) {
			case READ_IP:
				strncpy(node_ip, buffer, i - start);
				//fprintf(stdout, "  Node IP set to %s\n", node_ip);
				break;
			case READ_OFFSET:
				if (sscanf(buffer, "%f", &offset_sec) != 1) {
					error = 1;
				}
				else {
					//fprintf(stdout, "  Offset set to %f\n", offset_sec);
				}
				break;
			case READ_TX_TIME:
				if (sscanf(buffer, "%lld", &time_transmitted_ns) != 1) {
					error = 1;
				}
				else {
					fprintf(stderr, "  TX time set to %lld\n",time_transmitted_ns);
				}
				break;
			case READ_ID:
				if (sscanf(buffer, "%lld", &id) != 1) {
					error = 1;
				}
				else {
					//fprintf(stderr, "  id time set to %lld\n",id);
				}
				break;
		}
		if (error) {
			fprintf(stderr, "  Error reading packet in state %d\n", state);
			break;
		}
		state++;
		//fprintf(stderr, "  State is %d\n", state);
		if (state == DONE) {
			break;
		}
		memset(buffer, 0, sizeof(buffer));
		start = i + 1;
  }

	

	// Filter out corrupted rows
	if (state != DONE) {
		printf("FAIL,WRONG_FORMAT,");
    return;
	}

	long long raw_time_transmitted = time_transmitted_ns;
	time_transmitted_ns += (long long)(1000000000.0 * offset_sec); 

	long long transit_time_ns = time_received_ns - time_transmitted_ns;
	if ((transit_time_ns / 1000000000ll) > 100) {
		printf("FAIL,TRANSIT_TOO_LONG,%f,%lld,%lld,%lld,%lld, received %s",offset_sec,raw_time_transmitted,time_transmitted_ns,time_received_ns,transit_time_ns,payload);
    return;
	}
 
  printf("%lld,%s,%lld,%lld,%lld,", id, node_ip, time_transmitted_ns, time_received_ns, transit_time_ns);
}

void print_usage() {
  fprintf(stderr, "Usage: rx local_ip [device]\n");
}

char* get_default_device() {
  char *dev;
  char errbuf[PCAP_ERRBUF_SIZE];
  dev = pcap_lookupdev(errbuf);
  if (dev == NULL) {
    fprintf(stderr, "Couldn't find default device: %s\nAre you root?\n", errbuf);
    return NULL;
  }
  return dev;
}

char* get_device(int argc, char** argv) {
  char* dev;
  if (argc > 2) {
    dev = argv[2];
  }
  else {
    dev = get_default_device();
  }
  if (dev == NULL) {
    fprintf(stderr, "No device selected. Quitting.\n");
    print_usage();
    exit(EXIT_FAILURE);
  }
  return dev;
}

char* get_local_ip(int argc, char** argv) {
  char* ip_str = NULL;
  if (argc > 1) {
    ip_str = argv[1];
  }
  if (ip_str == NULL) {
    fprintf(stderr, "No local IP specified. Quitting.\n");
    print_usage();
    exit(EXIT_FAILURE);
  }
  // 255.255.255.255\0
  char* filtered_ip_str = malloc(16*sizeof(char));
  strncpy(filtered_ip_str, ip_str, 15);
  filtered_ip_str[15] = 0;
  return filtered_ip_str;
}

volatile sig_atomic_t should_update_ntp_offset = 0;

void get_ntp_offset() {
	should_update_ntp_offset = 0;
	char offset_buffer[30];
	FILE* offset_file = fopen("receiver_offset", "r");	
	fgets(offset_buffer, sizeof(offset_buffer), offset_file);
	fclose(offset_file);
	//fprintf(stderr, "Offset file contents: %s\n", offset_buffer);
	float time_offset;
  sscanf(offset_buffer, "%f", &time_offset);
  ns_offset = time_offset * 1000000000;
	fprintf(stderr, "Set NTP offset to %lld ns (%f s)\n", ns_offset, time_offset);	
}

void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{

  static int count = 1;                   /* packet counter */
  
  /* declare pointers to packet headers */
  //const struct sniff_ethernet *ethernet;  /* The ethernet header [1] */
  const struct sniff_ip *ip;              /* The IP header */
  //const struct tcphdr *tcp;            /* The TCP header */
  const struct udphdr *udp;            /* The TCP header */
  const u_char *payload;                    /* Packet payload */
  long long time_received_ns;
  int size_ip;
  //int size_tcp;
  int size_payload;
  
  // received: (long long) time: ns + (float) offset: seconds
  
  // Real time received, ns
  time_received_ns = get_time_ns();
	
	if (should_update_ntp_offset) {
		get_ntp_offset();
	}

	time_received_ns += ns_offset;
	//fprintf(stderr, "NTP offset is %lld\n", ns_offset);

  //printf("\nPacket number %d:\n", count);
  count++;
  
  /* define ethernet header */
  //ethernet = (struct sniff_ethernet*)(packet);
  
  /* define/compute ip header offset */
  ip = (struct sniff_ip*)(packet + SIZE_ETHERNET);
  size_ip = IP_HL(ip)*4;
  if (size_ip < 20) {
    fprintf(stderr, "   * Invalid IP header length: %u bytes\n", size_ip);
    return;
  }

  switch(ip->ip_p) {
    case IPPROTO_TCP:
			break;
/*
      printf("%s,TCP,", inet_ntoa(ip->ip_src));
      
      tcp = (struct tcphdr*)(packet + SIZE_ETHERNET + size_ip);
      size_tcp = tcp->th_off*4;
      if (size_tcp < 20) {
        printf("   * Invalid TCP header length: %u bytes\n", size_tcp);
        return;
      }
      printf("%d,%d,", ntohs(tcp->th_sport), ntohs(tcp->th_dport));
      payload = (u_char *)(packet + SIZE_ETHERNET + size_ip + size_tcp);
      size_payload = ntohs(ip->ip_len) - (size_ip + size_tcp);
      
      if (size_payload > 0) {
				fprintf(stderr, "  %s\n", payload);
        print_payload(payload, size_payload, time_received_ns);
      }
      break;*/
    case IPPROTO_UDP:
			printf("%s,", inet_ntoa(ip->ip_src));
      
      udp = (struct udphdr*)(packet + SIZE_ETHERNET + size_ip);
      printf("%d,%d,", ntohs(udp->uh_sport), ntohs(udp->uh_dport));
      payload = (u_char*)(packet + SIZE_ETHERNET + size_ip + sizeof(struct udphdr));
      size_payload = ntohs(ip->ip_len) - (size_ip + sizeof(struct udphdr));

      if (size_payload > 0) {
        print_payload(payload, size_payload, time_received_ns);
      }
			printf("\n");
      break;
    case IPPROTO_ICMP:
      //printf("ICMP,");
      break;
    case IPPROTO_IP:
      //printf("IP,");
      break;
    default:
      //printf("UNKNOWN");
      break;
  }
}

void do_ntp_sync(int sig) {
	should_update_ntp_offset = 1;
}

void init_pcap_capture(char* dev, char* local_ip) {
  char errbuf[PCAP_ERRBUF_SIZE];
  pcap_t* handle;
  char filter_expression[32];
  bpf_u_int32 net;
  bpf_u_int32 mask;
  struct bpf_program filter_program;
  sprintf(filter_expression, "dst host %s and ip and not (src port 22 or dst port 22)", local_ip);
  if ((handle = pcap_open_live(dev, 1514, 1, 500, errbuf)) == NULL) {
    fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
    exit(EXIT_FAILURE);
  }
  if (pcap_lookupnet(dev, &net, &mask, errbuf) == -1) {
    fprintf(stderr, "Couldn't get netmask for device %s: %s\n", dev, errbuf);
    exit(EXIT_FAILURE);
  }
  if (pcap_compile(handle, &filter_program, filter_expression, 0, net) == -1) {
    fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_expression, pcap_geterr(handle));
    exit(EXIT_FAILURE);
  }
  if (pcap_setfilter(handle, &filter_program) == -1) {
    fprintf(stderr, "Couldn't install filter %s: %s\n", filter_expression, pcap_geterr(handle));
    exit(EXIT_FAILURE);
  }
  
	signal(SIGUSR1, do_ntp_sync);

  fprintf(stderr, "Starting capture, NTP offset is %lld ns\n", ns_offset);
  
  pcap_loop(handle, -1, got_packet, NULL);
  
  // Cleanup
  pcap_freecode(&filter_program);
  pcap_close(handle);
}

// usage: rx ip [device]
// NTP direction: negative values mean local clock is fast
// aka offset is distance of local clock from true time
int main(int argc, char** argv) {
  char* dev = get_device(argc, argv); // 2
  if (argc < 2) {
    print_usage();
    exit(EXIT_FAILURE);
  }
	get_ntp_offset();
  char* local_ip = get_local_ip(argc, argv); // 1
  //printf("The device is %s\n", dev);
  init_pcap_capture(dev, local_ip);
  return 0;
}
